#include "HomogeneousMedium.cuh"

using namespace OBR;

static __forceinline__ __device__ Spectrum calcTr(RayData &ray, HomogeneousMediumData *data)
{
    return (-data->sigma_t * ray.t).exp();
}

static __forceinline__ __device__ Spectrum trackInteraction(SamplePRD *prd,
                                                            HomogeneousMediumData *data)
{
    float4 urand = curand_uniform4(prd->rng_long);

    int ch = urand.x * Spectrum::nSamples;
    ch = clamp(ch, 0, Spectrum::nSamples - 1);
    float dist = -logf(1 - urand.y) / data->sigma_t[ch];

    bool exist_mi = dist < prd->inter.wo.t;
    RayData &wo = prd->inter.wo;
    if (dist < prd->inter.wo.t)
    {
        /* ----------------- process ray & interaction data ----------------- */
        InteractionData &mi = prd->inter;
        wo.t = dist;
        mi.type = InteractionType::MEDIUM_INTERACTION;
        mi.world_pos = wo.o + wo.d * dist;

        /* ---------------------------- sample wi --------------------------- */
        float g = data->g;
        // Compute $\cos \theta$ for Henyey--Greenstein sample
        float cosTheta;
        if (fabs(g) < 1e-3)
        {
            cosTheta = 1 - 2 * urand.z;
        }
        else
        {
            float sqrTerm = (1 - g * g) / (1 + g - 2 * g * urand.z);
            cosTheta = -(1 + g * g - sqrTerm * sqrTerm) / (2 * g);
        }
        // Compute direction _wi_ for Henyey--Greenstein sample
        float sinTheta = sqrtf(fmaxf((float)0, 1 - cosTheta * cosTheta));
        float phi = 2 * M_PIf * urand.w;
        // generate new ray
        RayData wi = RayData();
        float3 v1, v2;
        coordinateSystem(wo.d, v1, v2);
        wi.d = sinTheta * cos(phi) * v1 + cosTheta * cos(phi) * v2 + sin(phi) * wo.d;
        wi.o = mi.world_pos;
        wi.has_differential = false;
        wi.med_id = wo.med_id;
        wi.med_instance_id = wo.med_instance_id;
        prd->inter.wi = wi;
    }

    // calculate pdf
    Spectrum tr = calcTr(wo, data);
    Spectrum density = exist_mi ? (data->sigma_t * tr) : tr;
    float pdf = 0;
    for (int i = 0; i < Spectrum::nSamples; ++i)
        pdf += density[i];
    pdf *= 1 / (float)Spectrum::nSamples;
    pdf = pdf == 0 ? 1 : pdf;

    // return beta
    return exist_mi ? (tr * data->sigma_s / pdf) : (tr / pdf);
}

extern "C" __device__ Spectrum __direct_callable__homogeneous_medium(SamplePRD *prd,
                                                                     MediumCallMode mode)
{
    // get sbt data
    CallableData *sbt_record = reinterpret_cast<CallableData *>(optixGetSbtDataPointer());
    HomogeneousMediumData *sbt_data =
        reinterpret_cast<HomogeneousMediumData *>(sbt_record->callable_payload);
    sbt_data += prd->inter.wo.med_instance_id;

    if (mode == MediumCallMode::TRACK_INTERACTION)
    {
        return trackInteraction(prd, sbt_data);
    }
    else if (mode == MediumCallMode::CALC_TR)
    {
        return calcTr(prd->inter.wo, sbt_data);
    }

    return Spectrum(0.0f);
}